package com.jboltai.capability.model.vllm;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ThinkStatus;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.util.http.HttpUtil;
import okhttp3.*;
import okhttp3.sse.EventSource;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

public class VLLMAIImpl implements AIEventProcessor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public static final VLLMAIImpl INSTANCE = new VLLMAIImpl();

    private static final String NORMAL_REQUEST = "NORMAL_REQUEST";
    private static final String STREAM_REQUEST = "STREAM_REQUEST";

    private static final String LOG_START = "[vLLM chat event][{}]: ";

    private static final String STOP_FLAG = "[DONE]";

    @Override
    public void process(AIChatEvent event, AIResource resource) {
        OkHttpClient client = resource.getClient();
        Request request = request(event, resource);

        LOGGER.debug("开始 调用Vllm大模型 处理AIChat事件：\n\n\nEvent参数: {}, \n\n\nRequest参数: {}", JSON.toJSONString(event), JSON.toJSONString(request));
        if (event.getIsStream()) {
            stream(event, resource, client, request);
        } else {
            normal(event, resource, client, request);
        }
    }

    private Request request(AIChatEvent event, AIResource resource) {
        String apiUrl = resource.getApiUrl();
        JSONObject param = param(event);
        RequestBody body = RequestBody.create(param.toJSONString(), MediaType.parse("application/json"));
        return new Request.Builder().url(apiUrl).post(body).header("Content-Type", "application/json").build();
    }

    private JSONObject param(AIChatEvent event) {
        JSONObject result = new JSONObject();
        result.put(AIParamKey.MODEL, event.getModelName());
        result.put(AIParamKey.STREAM, event.getIsStream());

        Double temperature = event.getTemperature();
        if (temperature != null) {
            result.put(AIParamKey.TEMPERATURE, temperature);
        }

        Double topP = event.getTopP();
        if (topP != null) {
            result.put(AIParamKey.TOP_P, topP);
        }

        Integer maxTokens = event.getMaxTokens();
        if (maxTokens != null) {
            result.put(AIParamKey.MAX_TOKENS, maxTokens);
        }

        JSONArray messages = new JSONArray(event.getMessages().size());
        result.put(AIParamKey.MESSAGES, messages);

        List<AIMessage> list = event.getMessages();
        for (AIMessage aiMessage : list) {
            JSONObject message = new JSONObject();
            switch (aiMessage.getRole()) {
                case user:
                    message.put("role", "user");
                    break;
                case assistant:
                    message.put("role", "assistant");
                    break;
                case system:
                    message.put("role", "system");
                    break;
                default:
                    continue;
            }
            message.put("content", aiMessage.getContent());
            messages.add(message);
        }
        return result;
    }

    private void normal(AIChatEvent event, AIResource resource, OkHttpClient client, Request request) {
        Call call = client.newCall(request); // 回调

        // 请求成功消费者
        Consumer<Response> successHandler = response -> {
            try {
                String jsonStr = response.body().string();
                LOGGER.info(LOG_START + "end process ... content: {}", event.getId(), jsonStr);

                if (response.code() != 200) {
                    handleResponseFail(event, resource, response.code(), jsonStr);
                    return;

                }

                AIMessage message = handleResponseSuccess(event, jsonStr, "message", false);
                event.setResult(message);
                event.setState(EventState.SUCCESS);
            } catch (Exception e) {
                LOGGER.error(LOG_START + "error: {}", event.getId(), e);
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                event.complete();
            }
        };

        // 请求失败消费者
        Consumer<Throwable> failHandler = throwable -> {
            LOGGER.error(LOG_START + "error: {}", event.getId(), throwable);
            event.fail(EventErrorType.NET_ERROR, throwable);
        };

        Callback callback = new Callback() {
            @Override
            public void onFailure(@NotNull Call call, @NotNull IOException e) {
                failHandler.accept(e);
            }

            @Override
            public void onResponse(@NotNull Call call, @NotNull Response response) throws IOException {
                try (Response r = response) {
                    successHandler.accept(r);
                }
            }
        };

        call.enqueue(callback);
    }

    private void stream(AIChatEvent event, AIResource resource, OkHttpClient client, Request request) {
        // 请求成功消费者
        Consumer<String> successHandler = body -> {
            try {
//                LOGGER.info(LOG_START + "end process ... content: {}", event.getId(), body);
                if (STOP_FLAG.equals(body)) {
                    event.setState(EventState.SUCCESS);
                    //结束
                    return;
                }
                handleResponseSuccess(event, body, "delta", true);
            } catch (Exception e) {
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                if (STOP_FLAG.equals(body)) {
                    event.complete();
                }
            }
        };
        // 请求失败消费者
        BiConsumer<Response, Throwable> failHandler = (response, throwable) -> {
            try {
                if (throwable != null) {
                    // 异常
                    handleException(event, throwable);
                    return;
                }
                if (response == null || response.body() == null) {
                    return;
                }
                // API返回错误
                String jsonStr = response.body().string();
                handleResponseFail(event, resource, response.code(), jsonStr);
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        };

        EventSource eventSource = HttpUtil.newEventSource(client, request, successHandler, failHandler);

        event.setAttr(STREAM_REQUEST, eventSource);
    }

    private AIMessage handleResponseSuccess(AIChatEvent event, String response,String contentField, boolean stream) {
        JSONObject data = JSON.parseObject(response);

        String id = data.getString("id");
        String content = "";
        if (data.containsKey("choices")) {
            JSONObject choice = data.getJSONArray("choices").getJSONObject(0);
            // 获取回复的消息
            content = choice.getJSONObject(contentField).getString("content");


            if (choice.containsKey("finish_reason")) {
                String finishReason = choice.getString("finish_reason");
                if(finishReason != null){
                    switch (finishReason) {
                        case "stop":
                            event.setFinishReason(MessageFinishReason.STOP);
                            break;
                        case "length":
                            event.setFinishReason(MessageFinishReason.LENGTH_LIMIT);
                            break;
                    }
                }
            }
        }

        AIMessage aiMessage = new AIMessage();
        aiMessage.setId(id);
        aiMessage.setRole(AIMessageRole.assistant);

        content = processThinking(event, aiMessage,  content, stream);

        aiMessage.setContent(content);

        event.setResultId(aiMessage.getId());
        event.setResultRole(aiMessage.getRole());
        event.appendResponse(aiMessage.getContent());

        // tokens
        if (data.containsKey("usage")) {
            JSONObject usage = data.getJSONObject("usage");
            event.setPromptTokenCount(usage.getInteger("prompt_tokens"));
            event.setTotalTokenCount(usage.getInteger("total_tokens"));
            event.setCompletionTokenCount(usage.getInteger("completion_tokens"));
        }

        if (StrUtil.isNotBlank(aiMessage.getContent())) {
            event.success(aiMessage);
        }
        return aiMessage;
    }

    /**
     * 处理深度思考内容
     * @param content
     * @param stream
     * @return 返回剔除深度思考后的ai回复内容，null代表目前都是深度思考内容，没有回复内容
     */
    private static String processThinking(AIChatEvent event, AIMessage message, String content, boolean stream) {
        if (stream == false) {
            //非流式
            if (content != null && content.contains("<think>") && content.contains("</think>" )) {
                //提取出思考内容
                int endIndex = content.lastIndexOf("</think>");
                String think = content.substring(content.indexOf("<think>") + 7, endIndex);
                event.thinkOver(think);
                message.setThink(think);
                return content.substring(endIndex + 8); //将深度思考之后的正常回复内容返回
            } else {
                return content;
            }

        }

        //流式
        if (content != null && content.contains("<think>")) {
            String think = content.substring(content.indexOf("<think>") + 7);
            event.thinking(think);
            message.setThink(think);
            return null;
        }

        if (content != null && content.contains("</think>" ) && event.getThinkStatus() == ThinkStatus.RUNNING) {
            String think = content.substring(0, content.indexOf("</think>") );
            event.thinkOver(think);
            message.setThink(think);
            return null;
        }

        if (event.getThinkStatus() == ThinkStatus.RUNNING || event.getThinkStatus() == ThinkStatus.START) {
            event.thinking(content);
            message.setThink(content);
            return null;
        }

        return content;


    }

    private static void handleException(AIEvent event, Throwable throwable) {
        LOGGER.error(LOG_START + "exception: {}", event.getId(), throwable);
        event.fail(EventErrorType.NET_ERROR, throwable);
    }

    private static void handleResponseFail(AIEvent event, AIResource resource, int code, String response) {
        JSONObject error;
        LOGGER.error(LOG_START + "response fail, code: {}, response: {}", event.getId(), code, response);
        switch (code) {
            case 400: {
                error = JSON.parseObject(response);
                event.fail(EventErrorType.REQUEST_PARAM_ERROR, error.getJSONObject("error").getString("message"));
                break;
            }

            case 401:
                event.fail(EventErrorType.API_KEY_INVALID, response);
                break;

            case 429: {
                error = JSON.parseObject(response);
                if ("rate_limit".equals(error.getJSONObject("error").getString("code"))) {
                    AIResourceCenter.moveResourceToInvalid(resource);
                    event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, response);
                } else {
                    event.fail(EventErrorType.ACCOUNT_BALANCE_INSUFFICIENT, response);
                }
                break;
            }

            case 500:
            case 503:
                event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                break;

            default:
                event.fail(EventErrorType.UNKNOWN, response);
                break;
        }
    }

    @Override
    public void handleCancel(AIChatEvent event, AIResource resource) {
        LOGGER.warn(LOG_START + "cancel ...", event.getId());

        Call call = event.getAttr(NORMAL_REQUEST);
        if (call != null) {
            call.cancel();
        }

        EventSource eventSource = event.getAttr(STREAM_REQUEST);
        if (eventSource != null) {
            eventSource.cancel();
        }
    }
}
